package cn.bincker.mybatis.encrypt.data.migration.impl;

import cn.bincker.mybatis.encrypt.data.migration.entity.MigrateTable;
import cn.bincker.mybatis.encrypt.data.migration.exception.MigrateException;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

import javax.sql.DataSource;
import java.io.Closeable;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collectors;

@Slf4j
public class InsertTask implements Runnable, Closeable {
    @Getter
    private boolean running = true;
    private Boolean destroyed = false;
    private final DataSource dataSource;
    private MigrateTable table;
    private final Function<MigrateTable, String> insertSqlBuilder;
    private final BiConsumer<Connection, MigrateTable> beforeInsert;
    private final BlockingQueue<Object[]> queue;
    private final int batchInsertRows;
    private final AtomicLong migratedSuccessRows;
    private final AtomicLong migratedFailedRows;
    private final boolean failContinue;
    private final Object busyLock = new Object();
    private transient boolean busy;

    public InsertTask(DataSource dataSource, Function<MigrateTable, String> insertSqlBuilder, BiConsumer<Connection, MigrateTable> beforeInsert, BlockingQueue<Object[]> queue, int batchInsertRows, AtomicLong migratedSuccessRows, AtomicLong migratedFailedRows, boolean failContinue) {
        this.dataSource = dataSource;
        this.insertSqlBuilder = insertSqlBuilder;
        this.beforeInsert = beforeInsert;
        this.queue = queue;
        this.batchInsertRows = batchInsertRows;
        this.migratedSuccessRows = migratedSuccessRows;
        this.migratedFailedRows = migratedFailedRows;
        this.failContinue = failContinue;
    }

    public void setMigrateTable(MigrateTable table){
        this.table = table;
    }

    @Override
    public void close() {
        running = false;
    }

    public void waitDestroyed() {
        if (destroyed) return;
        synchronized (this){
            try {
                this.wait();
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
    }

    @Override
    public void run() {
        var data = new ArrayList<Object[]>(batchInsertRows);
        try (var conn = dataSource.getConnection()){
            conn.setAutoCommit(false);
            while (running){
                queue.drainTo(data, batchInsertRows);
                if (data.isEmpty()){
                    synchronized (busyLock){
                        busy = false;
                        busyLock.notifyAll();
                    }
                    continue;
                }
                busy = true;

                beforeInsert.accept(conn, table);
                String sql = insertSqlBuilder.apply(table);
                try(var insertStatement = conn.prepareStatement(sql)){
                    for (Object[] row : data) {
                        if (row == null) continue;
                        for (int j = 0; j < row.length; j++) {
                            try {
                                insertStatement.setObject(j + 1, row[j]);
                            }catch (SQLException e){
                                log.error("set params fail: index={}\tsql={}", j + 1, sql);
                                throw e;
                            }
                        }
                        insertStatement.addBatch();
                    }
                    var executeResult = insertStatement.executeBatch();
                    conn.commit();
                    table.getMigratedRows().addAndGet(executeResult.length);
                    migratedSuccessRows.getAndAdd(Arrays.stream(executeResult).sum());
                    long failedCount = Arrays.stream(executeResult).filter(i -> i == 0).count();
                    migratedFailedRows.getAndAdd(failedCount);
                    for (int i = 0; i < executeResult.length; i++) {
                        if (executeResult[i] == 0){
                            int finalI = i;
                            log.error(
                                    "row insert fail: table={}\tkey={}",
                                    table.getSrcTableName(),
                                    table.getKeys().stream().map(k->data.get(finalI)[k.getIndex()].toString()).collect(Collectors.joining(","))
                            );
                        }
                    }
                    if (!failContinue && failedCount > 0){
                        throw new MigrateException("insert fail: rows=" + failedCount);
                    }
                }
                data.clear();

                if(queue.isEmpty()) {
                    synchronized (queue) {
                        queue.notifyAll();
                    }
                }
                if (table.getMigratedRows().get() >= table.getTotalRows()){
                    //noinspection SynchronizeOnNonFinalField
                    synchronized (table){
                        table.notifyAll();
                    }
                }
            }
        } catch (SQLException e) {
            log.error("insert data error", e);
        } finally {
            running = false;
        }
        destroyed = true;
        synchronized (this) {
            this.notifyAll();
        }
    }

    public void waitBusy() {
        while (busy){
            synchronized (busyLock) {
                if (!busy) return;
                try {
                    busyLock.wait(1000);
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }
            }
        }
    }
}
